﻿#include "sky_seger.h"

#include <math.h>

#include <thread>

#include "timer_utils.h"

#include <xtensor/xadapt.hpp>
#include <xtensor/xarray.hpp>
#include <xtensor/xshape.hpp>
#include <xtensor/xio.hpp>

using namespace std;
using namespace xt;

const std::string SkySeger::sky_seg_model_path = "/system/model/sky_ground_model.rknn";
rknn_context SkySeger::template_session_ = 0;
void* SkySeger::template_modle_ = nullptr;

int SkySeger::input_height_ = 0;
int SkySeger::input_width_ = 0;
int SkySeger::output_height_ = 0;
int SkySeger::output_width_ = 0;

bool SkySeger::model_loaded = 0;

int SkySeger::start(const cv::Mat& img, cv::Mat& img_sky_mask) {
    time_checker timer;

    if (model_loaded == 0) {
        if (loadModel() < 0)
            return -1;
    }

    // 预处理
    timer.start();

    cv::Mat img_input;
    int ret = preprocess(img, img_input);
    if (ret < 0)
        return ret;

    timer.stop();
    timer.show_distance("preprocess cost");

    // NPU
    timer.start();

    vector<float> output_data;
    process(img_input, output_data);

    timer.stop();
    timer.show_distance("process cost");

    // 后处理
    timer.start();

    postprocess(output_data, img_sky_mask);

    timer.stop();
    timer.show_distance("postprocess cost");

    return 0;
}

int SkySeger::preprocess(const cv::Mat& img, cv::Mat& img_input) {
    if (img.empty()) {
        std::cout << "read failed"  << std::endl;
        return -1;
    }

    cv::resize(img, img_input, cv::Size(input_width_, input_height_)); // 调整图像尺寸
    cv::cvtColor(img_input, img_input, cv::COLOR_BGR2RGB);

    return 0;
}

int SkySeger::process(const cv::Mat& img_input, vector<float>& output_data) {
    vector<rknn_input> inputs(1);
    memset(inputs.data(), 0, inputs.size());
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = img_input.rows * img_input.cols * img_input.channels();
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    inputs[0].buf = img_input.data;

    rknn_inputs_set(template_session_, 1, inputs.data());

    vector<rknn_output> outputs(1);
    memset(outputs.data(), 0, outputs.size());
    outputs[0].want_float = 1;

    rknn_run(template_session_, 0);
    rknn_outputs_get(template_session_, 1, outputs.data(), NULL);

    output_data.assign(static_cast<float*>(outputs[0].buf), static_cast<float*>(outputs[0].buf) + output_width_ * output_height_);

    rknn_outputs_release(template_session_, 1, outputs.data());

    return 0;
}

int SkySeger::postprocess(const vector<float>& output_data, cv::Mat& img_sky_mask) {
    std::vector<float> output_sigmoid;
    output_sigmoid.clear();

    for (int i = 0; i < output_width_ * output_height_; i++)
        output_sigmoid.push_back(sigmoid(output_data[i]));

    // 将概率预测值转换为二值掩码图像
    float threshold = 0.5;
    img_sky_mask = cv::Mat(cv::Size(output_width_, output_height_), CV_8UC1);

    for (int y = 0; y < output_height_; ++y) {
        for (int x = 0; x < output_width_; ++x) {
            float value = output_sigmoid[y * output_width_ + x];
            img_sky_mask.at<uchar>(y, x) = (value > threshold) ? 255 : 0;
        }
    }

    cv::imwrite("/tmp/sky_seger.jpg", img_sky_mask);

    return 0;
}

int SkySeger::loadModel() {
    if (model_loaded == 0) {
        FILE* template_fp = fopen(sky_seg_model_path.c_str(), "rb");
        if (template_fp == nullptr)
            return -1;

        fseek(template_fp, 0, SEEK_END);
        long template_fsize = ftell(template_fp);
        template_modle_ = malloc(template_fsize);
        fseek(template_fp, 0, SEEK_SET);
        size_t template_len = fread((unsigned char*)template_modle_, 1, template_fsize, template_fp);

        int ret = rknn_init(&template_session_, template_modle_, template_fsize, 0);
        if (ret < 0) {
            log_e("rknn_init error ret=%d", ret);
            return -1;
        }

        printRKNNAttr(template_session_);

        fclose(template_fp);

        model_loaded = 1;
    } else {
        log_i("sky seger model loaded");
    }

    return 0;
}

// 有误差导致分割结果与python不同
float SkySeger::fastExp(float x) {
    union {
        uint32_t i;
        float f;
    } v{};
    v.i = (1 << 23) * (1.4426950409 * x + 126.93490512f);
    return v.f;
}

float SkySeger::sigmoid(float x) {
    // return 1.0f / (1.0f + fastExp(-x));
    return 1.0 / (1.0 + expf(-x));
}

void SkySeger::printRKNNTensor(rknn_tensor_attr *attr) {
    log_i("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
        "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f",
        attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
        attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
        attr->qnt_type, attr->fl, attr->zp, attr->scale);
}

int SkySeger::printRKNNAttr(rknn_context ctx) {
    rknn_input_output_num io_num;
    int ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret < 0) {
        log_i("rknn_init error ret=%d", ret);
        return -1;
    }
    int input_num = io_num.n_input;
    int output_num = io_num.n_output;
    log_i("model input num: %d, output num: %d", input_num, output_num);

    rknn_tensor_attr input_attrs[input_num];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < input_num; i++) {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]),
                            sizeof(rknn_tensor_attr));
        printRKNNTensor(&(input_attrs[i]));

        // NHWC
        input_height_ = input_attrs[i].dims[2];
        input_width_ = input_attrs[i].dims[1];
    }

    rknn_tensor_attr output_attrs[output_num];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int k = 0; k < output_num; k++) {
        output_attrs[k].index = k;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[k]),
                            sizeof(rknn_tensor_attr));
        printRKNNTensor(&(output_attrs[k]));

        // NCHW
        output_height_ = output_attrs[k].dims[1];
        output_width_ = output_attrs[k].dims[0];
    }

    return 0;
}
